package cn.ichiva.dl.tushare.analy;

import cn.ichiva.dl.tushare.analy.impl.Statistical;
import cn.ichiva.dl.tushare.analy.impl.Trade;
import cn.ichiva.dl.tushare.analy.impl.TradeAgent;
import lombok.extern.slf4j.Slf4j;
import org.deeplearning4j.nn.multilayer.MultiLayerNetwork;
import org.nd4j.linalg.api.ndarray.INDArray;

import java.util.Arrays;
import java.util.List;

/**
 * Created by game group on 2020/12/8.
 */
@Slf4j
public class TrainHelper {

    public static boolean contend(TradeAgent agent) {
        partTrain(agent, P4Data.TRAIN_1);
        if(agent.getScores().get(0) < agent.getTrade().getInitMoney() * 1.1){
            log.debug("第一阶段 亏损丢弃 {}",agent);
            return false;
        }

        partTrain(agent, P4Data.TRAIN_2);
        if(agent.getScores().get(1) < agent.getTrade().getInitMoney() * 1.1){
            log.debug("第二阶段 亏损丢弃 {}",agent);
            return false;
        }

        //第三阶段数据不能参与训练
        test(agent);
        Statistical.writeAgent(agent);
        return true;
    }

    public static void partTrain(TradeAgent agent, List<DayINDArray> trainList) {
        Trade trade = agent.getTrade();
        trade.setMoney(trade.getInitMoney());
        trade.setNStock(0);

        MultiLayerNetwork network = agent.getMultiLayerNetwork();
        makingList(agent,network, trainList);
        double price =  trainList.get(trainList.size() -1).getData().getDouble(0,0,0);
        double stockMoney = agent.getTrade().getNStock() * price;
        agent.getScores().add(agent.getTrade().getMoney() + stockMoney);
    }

    private static void makingList(TradeAgent agent, MultiLayerNetwork network, List<DayINDArray> dataList) {
        int len = dataList.size() - 1;
        for (int i = 0; i < len; i++) {
            INDArray data = dataList.get(i).getData();
            INDArray output = network.rnnTimeStep(data);
            List<Double> list = Arrays.asList(output.getDouble(new long[]{0, 0,0}),
                    output.getDouble(new long[]{0, 1,0}),
                    output.getDouble(new long[]{0, 2,0}));
            double max = list.stream().mapToDouble(e -> e).max().getAsDouble();
            //使用第二天的开盘价
            DayINDArray p3Data = dataList.get(i + 1);
            double price = p3Data.getData().getDouble(0, 0,0);
            if(max == list.get(0)){
                agent.getTrade().buy(p3Data.getDay(),price);
            }else if(max == list.get(2)){
                agent.getTrade().sell(p3Data.getDay(),price);
            }else{
                agent.getTrade().waitAndSee(p3Data.getDay(),price);
            }
        }
    }

    public static void test(TradeAgent agent) {
        Trade trade = agent.getTrade();
        trade.setMoney(trade.getInitMoney());
        trade.setNStock(0);

        log.debug("验证... {}",agent);

        List<DayINDArray> testList = P4Data.TEST_0;
        MultiLayerNetwork network = agent.getMultiLayerNetwork();
        makingList(agent, network, testList);
        double price = testList.get(testList.size() - 1).getData().getDouble(0,0, 0);
        double stockMoney = agent.getTrade().getNStock() * price;
        agent.getScores().add(agent.getTrade().getMoney() + stockMoney);
    }
}
